package com.qiniuyun.backai.util;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.qiniuyun.backai.model.pojo.Chat;
import com.qiniuyun.backai.service.ChatService;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

//@Component
public class InSqlChatMemory implements ChatMemory {
    Map<String, List<Message>> conversationHistory = new ConcurrentHashMap<>();
    @Autowired
    private ChatService chatService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        //将传入的数据转为pojo
        Chat chat = new Chat();
        chat.setConvid(Long.valueOf(conversationId));
        //判断是用户还是ai
        MessageType messageType = messages.get(0).getMessageType();
        if (messageType.equals(MessageType.USER)) {
            chat.setRole(0);
        } else {
            chat.setRole(1);
        }
        chat.setContenttype(0);
        chat.setAudiourl("");
        chat.setCreatetime(new Date());
        chat.setUpdatetime(new Date());
        //判断是否存在历史
        List<Message> history = get(conversationId, Integer.MAX_VALUE);
        Boolean isok = true;
        if (history.size() == 0) {
            isok = false;
        }
        this.conversationHistory.putIfAbsent(conversationId, new ArrayList<>());
        this.conversationHistory.get(conversationId).addAll(messages);
        JSONObject json = new JSONObject(conversationHistory);
        chat.setContent(json);
        //如果id存在就更新，反之就保存
        if (isok) {
            chatService.updateByConvId(chat, conversationId);
        } else {
            chatService.save(chat);
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        //从当前缓存中获取
        List<Message> messages = this.conversationHistory.get(conversationId);
        Chat chat = new Chat();
        List<Message> all = new ArrayList<>();
        //如果缓存中不存在从数据库中缓存
        if (messages == null) {
            chat = chatService.getByConvId(String.valueOf(conversationId));

            if (chat != null) {
                Map<String, Object> map = JSON.parseObject((String) chat.getContent(), Map.class);
                List<JSONObject> jsonMessages = (List<JSONObject>) map.get(conversationId);
                if (jsonMessages != null) {
                    all = jsonMessages.stream()
                            .map(json -> {
                                String text = String.valueOf(json.get("text"));
                                String role = String.valueOf(json.get("messageType"));
                                if ("USER".equals(role)) {
                                    return new UserMessage(text);
                                } else if ("ASSISTANT".equals(role)) {
                                    return new AssistantMessage(text);
                                }
                                return null; // 其他情况可以扩展
                            })
                            .filter(Objects::nonNull)
                            .collect(Collectors.toList());
                }
            }
            //放入缓存
            this.conversationHistory.put(conversationId, all);
        } else {
            return messages;
        }
        return all != null ? all.stream().skip(Math.max(0, all.size() - lastN)).toList() : List.of();
    }

    @Override
    public void clear(String conversationId) {
        this.conversationHistory.remove(conversationId);
    }
}
